import vtk
import numpy as np

# ----------------------------
# 1. 准备示例数据（替换为你的x,y,z,u,v,w）
# ----------------------------
# 生成随机离散点（示例：1000个点）
np.random.seed(42)
n_points = 1000
x = np.random.uniform(-2, 2, n_points)  # x坐标
y = np.random.uniform(-2, 2, n_points)  # y坐标
z = np.random.uniform(-2, 2, n_points)  # z坐标

# 定义速度场（示例：旋转+轴向流动）
u = y  # x方向速度
v = -x  # y方向速度
w = z  # z方向速度

# ----------------------------
# 2. 将数据导入VTK数据结构
# ----------------------------
# 创建点集
points = vtk.vtkPoints()
for i in range(n_points):
    points.InsertNextPoint(x[i], y[i], z[i])

# 创建速度向量数组
vectors = vtk.vtkDoubleArray()
vectors.SetName("Velocity")
vectors.SetNumberOfComponents(3)  # u, v, w
for i in range(n_points):
    vectors.InsertNextTuple3(u[i], v[i], w[i])

# 封装为vtkPolyData
point_data = vtk.vtkPolyData()
point_data.SetPoints(points)
point_data.GetPointData().SetVectors(vectors)

# ----------------------------
# 3. 插值离散点为连续网格场（关键步骤）
# ----------------------------
# 方法：用高斯平滑将离散点插值到规则网格
splatter = vtk.vtkGaussianSplatter()
splatter.SetInputData(point_data)
splatter.SetSampleDimensions(50, 50, 50)  # 插值网格分辨率
splatter.SetRadius(0.5)  # 高斯核半径（控制平滑程度）
splatter.SetExponentFactor(-10)  # 高斯衰减因子
# splatter.SetVectorArrayName("Velocity")  # 要插值的向量场名称
splatter.Update()

# 提取插值后的网格数据（vtkImageData）
volume = splatter.GetOutput()
print(f"高斯散射完成，体积数据维度: {volume.GetDimensions()}")

# ----------------------------
# 4. 生成流线（vtkStreamTracer）
# ----------------------------
# 定义种子点源（例如：在z=0平面生成网格状种子）
seed_source = vtk.vtkPlaneSource()
seed_source.SetOrigin(-2, -2, 0)
seed_source.SetPoint1(2, -2, 0)
seed_source.SetPoint2(-2, 2, 0)
seed_source.SetXResolution(8)  # x方向种子数
seed_source.SetYResolution(8)  # y方向种子数

# 流线追踪器
stream_tracer = vtk.vtkStreamTracer()
stream_tracer.SetInputData(volume)  # 输入插值后的速度场
stream_tracer.SetSourceConnection(seed_source.GetOutputPort())  # 种子点
stream_tracer.SetMaximumPropagation(200)  # 最大积分步数
stream_tracer.SetInitialIntegrationStep(0.1)  # 初始步长
stream_tracer.SetIntegrationDirectionToBoth()  # 双向积分
stream_tracer.Update()

streamlines = stream_tracer.GetOutput()
print(f"流线追踪完成，生成 {streamlines.GetNumberOfLines()} 条流线")
print("系统出现")

# ----------------------------
# 5. 可视化设置
# ----------------------------
# 流线映射器
stream_mapper = vtk.vtkPolyDataMapper()
stream_mapper.SetInputConnection(stream_tracer.GetOutputPort())
# 用速度大小着色
stream_mapper.SetScalarModeToMagnitude()
stream_mapper.SetScalarRange(volume.GetPointData().GetVectors().GetRange(-1))  # 颜色范围

# 流线演员
stream_actor = vtk.vtkActor()
stream_actor.SetMapper(stream_mapper)
stream_actor.GetProperty().SetLineWidth(2)

# 颜色条
color_bar = vtk.vtkScalarBarActor()
color_bar.SetLookupTable(stream_mapper.GetLookupTable())
color_bar.SetTitle("Velocity Magnitude")

# ----------------------------
# 6. 渲染与显示
# ----------------------------
renderer = vtk.vtkRenderer()
render_window = vtk.vtkRenderWindow()
render_window.AddRenderer(renderer)
interactor = vtk.vtkRenderWindowInteractor()
interactor.SetRenderWindow(render_window)

renderer.AddActor(stream_actor)
renderer.AddActor(color_bar)
renderer.SetBackground(0.1, 0.1, 0.1)
renderer.ResetCamera()

render_window.SetSize(800, 600)
render_window.SetWindowName("Streamlines from x,y,z,u,v,w")
render_window.Render()
interactor.Start()